import cv2
import gradio as gr     
import numpy as np
from PIL import Image as Image 
from sklearn.cluster import KMeans

def analyze_palette(image):  # 定义 gradio 应用函数
    img_rgb = image
    pixels = img_rgb.reshape(-1, 3)  # 重塑像素数组
    
    kmeans = KMeans(n_clusters=5)          # 创建 KMeans 模型
    kmeans.fit_predict(pixels)           # 拟合模型
    colors_rgb = kmeans.cluster_centers_        # 获取聚类中心
    colors_hex = ['#' + ''.join(f'{int(c):02x}' for c in color) for color in colors_rgb]    # 转换为十六进制码

    # 创建颜色块图像
    palette = np.zeros((50, 250, 3), dtype=np.uint8)    # 创建空白图像
    steps = palette.shape[1] // kmeans.n_clusters   # 计算每个颜色块的宽度
    for i, color in enumerate(colors_rgb):      # 为每个颜色块填充颜色，enumerate() 函数用于将一个可遍历的数据对象组合为一个索引序列
        palette[:, i*steps:(i+1)*steps] = color     # 填充颜色, i*steps:(i+1)*steps 为每个颜色块的宽度, color 为颜色值

    # 将numpy数组转换为PIL图像
    palette_image = Image.fromarray(palette)

    # 返回颜色块图像和十六进制的颜色值
    return palette_image, "\n".join(colors_hex)

# 定义 Gradio 界面
iface = gr.Interface(
    fn=analyze_palette,
    inputs=gr.Image(type="numpy", label="上传图片"),
    outputs=[
        gr.Image(type="pil", label="颜色块"),  # 输出类型为 PIL Image
        gr.Text(label="色彩十六进制码")  # 输出类型为文本
    ],
    title="KMeans调色板提取器",
    description="上传图片来提取它的调色板。"
)

# 运行 Gradio 应用
iface.launch()